{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "05681fff-8d6a-48a3-a38f-ec6cab8b0418",
   "metadata": {},
   "source": [
    "# SplitRec：在隐语中使用拆分 BST 算法（Torch 后端）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4290e552-9930-40b1-bab8-4f24f3747608",
   "metadata": {},
   "source": [
    "阿里搜索推荐团队在 2019 年 arXiv 上发布文章《Behavior Sequence Transformer for E-commerce Recommendation in Alibaba》，提出 BST 模型。BST 利用 Transformer 结构捕捉用户行为序列信息，解决 WDL、DIN 等模型忽略序列信息的问题，因而受到关注。本文将介绍如何在隐语中使用拆分 BST 算法。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49c3da26-a348-4c38-b7ec-1389c9ee4309",
   "metadata": {},
   "source": [
    "## BST 模型\n",
    "BST 模型在 DNN 模型基础上，使用 Transformer Layer 捕捉用户行为序列中的信息，整体结构如下\n",
    "![bst](./resources/bst0.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "797c790a-4c3c-4e03-befe-16c4e3db14db",
   "metadata": {},
   "source": [
    "其中 Transformer Layer 使用一层 multi-head self-attention 结构，如图右侧所示。\n",
    "\n",
    "Transformer layer 的输入是序列特征的 Embedding 和位置 Embedding 的加和。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a925168a-d700-4fd7-b3c4-73b4d8908a0a",
   "metadata": {},
   "source": [
    "## 隐语中的拆分 BST 模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0242c4f-a72b-4a92-9736-38971146f27f",
   "metadata": {},
   "source": [
    "隐语考虑两方拆分学习场景，两方各持有一些特征，包括序列特征和其他特，一方持有 label 。双方特征经过 base 模型处理后，将 base 模型的输出输入到 fuse 模型进行学习，根据对序列特征的处理不同分为拆分 BST 模型和拆分 BSTPlus 模型。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6f7a22-4e9a-4a50-aac4-b1071fb035ff",
   "metadata": {},
   "source": [
    "### 拆分 BST 模型\n",
    "其中序列特征通过 Transformer Layer 处理，其他特征通过普通的 Embedding 处理，序列特征和其他特征是可选的，模型结构如图所示。\n",
    "\n",
    "![bst](./resources/bst1.png)\n",
    "\n",
    "### 拆分 BSTPlus 模型\n",
    "上述 BST 模型虽然能够有效捕捉序列信息，但对于同一用户不同商品的预估来说，序列信息都是相同的，并不能区分同一用户行为序列与不同商品的关联性，因此 BSTPlus 模型在 BST 基础上叠加 DIN 中的 attention 思想，对 target item 和 Transformer layer 的输出进行 attention pooling，通过计算 target item 对序列中 item 的注意力来挖掘用户行为序列与 target item 之间的关联信息，结构如下图。由于 target item 存在于 label 方，暂不支持无 label 方有序列特征。\n",
    "\n",
    "![bst](./resources/bst2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4096d2d3-8991-454a-8b67-53239860248e",
   "metadata": {},
   "source": [
    "## 隐语封装\n",
    "我们在隐语中提供了对于各种应用的封装。 关于 BST 的封装在 secretflow/ml/nn/applications/sl_bst_torch.py，提供了 `BSTBase` `BSTPlusBase` 和 `BSTFuse` 几个类。\n",
    "下面我们通过一个例子来看一下如何使用隐语封装的 BST 来进行训练。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9955b23-eb62-4dee-b88d-be85ba50a154",
   "metadata": {},
   "source": [
    "## 环境设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "66f50a5c-563f-4fd8-a263-b480283da87d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.1.0.dev20230926\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-09-26 19:48:53,161\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob'], address=\"local\", log_to_driver=False)\n",
    "alice, bob = sf.PYU('alice'), sf.PYU('bob')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba03e5f6-b537-4bdc-adef-4b4f2ba7267b",
   "metadata": {},
   "source": [
    "## 数据集介绍\n",
    "这里将使用最经典的 MovieLens 数据集进行演示。 MovieLens 是一个开放式的推荐系统数据集，包含了电影评分和电影元数据信息。\n",
    "\n",
    "[数据集官网](https://grouplens.org/datasets/movielens/)\n",
    "\n",
    "[下载数据集](https://secretflow-data.oss-accelerate.aliyuncs.com/datasets/movielens/ml-1m.zip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6726bbb9-a165-44fb-a6ac-6de652c0ff7b",
   "metadata": {},
   "source": [
    "## 下载并处理数据\n",
    "我们通过聚合同一用户按时间顺序评分过的电影生成用户行为序列，并对数据进行切分：  \n",
    "- alice: user_id, target_movie_id, sequence_movie_ids, label  \n",
    "- bob: gender, age_group, occupation  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "80a5cdb4-39c9-41f9-9613-7aaf5bb9b5bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "data_dir = './test_bst_data'\n",
    "dataset_download_dir = data_dir + '/data_download'\n",
    "gen_data_path = data_dir + '/data_sl_bst'\n",
    "fea_emb_input_size = {}\n",
    "\n",
    "\n",
    "def generate_data():\n",
    "    import numpy as np\n",
    "    import pandas as pd\n",
    "\n",
    "    from secretflow.utils.simulation.datasets import _DATASETS, get_dataset, unzip\n",
    "\n",
    "    global dataset_download_dir\n",
    "    # download dataset\n",
    "    if not Path(dataset_download_dir).is_dir():\n",
    "        filepath = get_dataset(_DATASETS['ml-1m'])\n",
    "        unzip(filepath, dataset_download_dir)\n",
    "\n",
    "    dataset_dir = dataset_download_dir + '/ml-1m'\n",
    "    # read data\n",
    "    users = pd.read_csv(\n",
    "        dataset_dir + \"/users.dat\",\n",
    "        sep=\"::\",\n",
    "        names=[\"user_id\", \"gender\", \"age_group\", \"occupation\", \"zip_code\"],\n",
    "        engine='python',\n",
    "    )\n",
    "\n",
    "    movies = pd.read_csv(\n",
    "        dataset_dir + \"/movies.dat\",\n",
    "        sep=\"::\",\n",
    "        names=[\"movie_id\", \"title\", \"genres\"],\n",
    "        engine='python',\n",
    "        encoding=\"ISO-8859-1\",\n",
    "    )\n",
    "\n",
    "    ratings = pd.read_csv(\n",
    "        dataset_dir + \"/ratings.dat\",\n",
    "        sep=\"::\",\n",
    "        names=[\"user_id\", \"movie_id\", \"rating\", \"unix_timestamp\"],\n",
    "        engine='python',\n",
    "    )\n",
    "\n",
    "    users[\"user_id\"] = users[\"user_id\"].apply(lambda x: f\"{x}\")\n",
    "    users[\"age_group\"] = users[\"age_group\"].apply(lambda x: f\"{x}\")\n",
    "    users[\"occupation\"] = users[\"occupation\"].apply(lambda x: f\"{x}\")\n",
    "\n",
    "    movies[\"movie_id\"] = movies[\"movie_id\"].apply(lambda x: f\"{x}\")\n",
    "    movies[\"genres\"] = movies[\"genres\"].apply(lambda x: ','.join(x.split('|')))\n",
    "\n",
    "    ratings[\"movie_id\"] = ratings[\"movie_id\"].apply(lambda x: f\"{x}\")\n",
    "    ratings[\"user_id\"] = ratings[\"user_id\"].apply(lambda x: f\"{x}\")\n",
    "    ratings[\"rating\"] = ratings[\"rating\"].apply(lambda x: float(x))\n",
    "\n",
    "    # encode movie_id\n",
    "    from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "    le = LabelEncoder()\n",
    "    le.fit(movies['movie_id'].unique())\n",
    "    movies['movie_id'] = (\n",
    "        le.transform(movies['movie_id']) + 1\n",
    "    )  # reserve index 0 for padding\n",
    "    movies['movie_id'] = movies['movie_id'].astype('string')\n",
    "    ratings[\"movie_id\"] = le.transform(ratings[\"movie_id\"]) + 1\n",
    "    ratings['movie_id'] = ratings['movie_id'].astype('string')\n",
    "    fea_emb_input_size['target_id'] = len(movies['movie_id'].unique()) + 1\n",
    "\n",
    "    # generate users' movie sequence\n",
    "    ratings_group = ratings.sort_values(by=[\"unix_timestamp\"]).groupby(\"user_id\")\n",
    "\n",
    "    ratings_data = pd.DataFrame(\n",
    "        data={\n",
    "            \"user_id\": list(ratings_group.groups.keys()),\n",
    "            \"movie_ids\": list(ratings_group.movie_id.apply(list)),\n",
    "            \"ratings\": list(ratings_group.rating.apply(list)),\n",
    "            \"timestamps\": list(ratings_group.unix_timestamp.apply(list)),\n",
    "        }\n",
    "    )\n",
    "\n",
    "    sequence_length = 5\n",
    "    step_size = 4\n",
    "\n",
    "    def create_sequences(values, window_size, step_size):\n",
    "        sequences = []\n",
    "        start_index = 0\n",
    "        while True:\n",
    "            end_index = start_index + window_size\n",
    "            seq = values[start_index:end_index]\n",
    "            if len(seq) < window_size:\n",
    "                if len(seq) > 1:\n",
    "                    seq.extend(['0'] * (window_size - len(seq)))\n",
    "                    sequences.append(seq)\n",
    "                break\n",
    "            sequences.append(seq)\n",
    "            start_index += step_size\n",
    "        return sequences\n",
    "\n",
    "    ratings_data.movie_ids = ratings_data.movie_ids.apply(\n",
    "        lambda ids: create_sequences(ids, sequence_length, step_size)\n",
    "    )\n",
    "\n",
    "    ratings_data.ratings = ratings_data.ratings.apply(\n",
    "        lambda ids: create_sequences(ids, sequence_length, step_size)\n",
    "    )\n",
    "    del ratings_data[\"timestamps\"]\n",
    "\n",
    "    ratings_data_movies = ratings_data[[\"user_id\", \"movie_ids\"]].explode(\n",
    "        \"movie_ids\", ignore_index=True\n",
    "    )\n",
    "    ratings_data_rating = ratings_data[[\"ratings\"]].explode(\n",
    "        \"ratings\", ignore_index=True\n",
    "    )\n",
    "    ratings_data_transformed = pd.concat(\n",
    "        [ratings_data_movies, ratings_data_rating], axis=1\n",
    "    )\n",
    "    ratings_data_transformed = ratings_data_transformed.join(\n",
    "        users.set_index(\"user_id\"), on=\"user_id\"\n",
    "    )\n",
    "    # last movie in sequence is target movie\n",
    "    ratings_data_transformed['movie_id'] = ratings_data_transformed.movie_ids.apply(\n",
    "        lambda x: x[-1] if '0' not in x else x[x.index('0') - 1]\n",
    "    )\n",
    "\n",
    "    # remove last movie from seqnuence to generate movie sequence\n",
    "    ratings_data_transformed.movie_ids = ratings_data_transformed.movie_ids.apply(\n",
    "        lambda x: (\n",
    "            \",\".join(x[:-1])\n",
    "            if '0' not in x\n",
    "            else \",\".join(x[: x.index('0') - 1] + x[x.index('0') :])\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # target movie's rating\n",
    "    ratings_data_transformed['label'] = ratings_data_transformed.ratings.apply(\n",
    "        lambda x: x[-1] if '0' not in x else x[x.index('0') - 1]\n",
    "    )\n",
    "\n",
    "    ratings_data_transformed.ratings = ratings_data_transformed.ratings.apply(\n",
    "        lambda x: \",\".join([str(v) for v in x[:-1]])\n",
    "    )\n",
    "    ratings_data_transformed = ratings_data_transformed.join(\n",
    "        movies.set_index(\"movie_id\"), on=\"movie_id\"\n",
    "    )\n",
    "\n",
    "    del (\n",
    "        ratings_data_transformed[\"zip_code\"],\n",
    "        ratings_data_transformed[\"title\"],\n",
    "        ratings_data_transformed[\"genres\"],\n",
    "        ratings_data_transformed[\"ratings\"],\n",
    "    )\n",
    "\n",
    "    ratings_data_transformed.rename(\n",
    "        columns={\"movie_ids\": \"sequence_movie_ids\", \"movie_id\": \"target_movie_id\"},\n",
    "        inplace=True,\n",
    "    )\n",
    "\n",
    "    # encoder: str to id\n",
    "    le = LabelEncoder()\n",
    "    ratings_data_transformed[['user_id', 'gender', 'age_group', 'occupation']] = (\n",
    "        ratings_data_transformed[\n",
    "            ['user_id', 'gender', 'age_group', 'occupation']\n",
    "        ].apply(le.fit_transform)\n",
    "    )\n",
    "    fea_emb_input_size['user_id'] = len(ratings_data_transformed['user_id'].unique())\n",
    "    fea_emb_input_size['gender'] = len(ratings_data_transformed['gender'].unique())\n",
    "    fea_emb_input_size['age_group'] = len(\n",
    "        ratings_data_transformed['age_group'].unique()\n",
    "    )\n",
    "    fea_emb_input_size['occupation'] = len(\n",
    "        ratings_data_transformed['occupation'].unique()\n",
    "    )\n",
    "\n",
    "    # split train/test\n",
    "    random_selection = np.random.rand(len(ratings_data_transformed.index)) <= 0.85\n",
    "    train_data = ratings_data_transformed[random_selection]\n",
    "    test_data = ratings_data_transformed[~random_selection]\n",
    "\n",
    "    if os.path.exists(gen_data_path):\n",
    "        shutil.rmtree(gen_data_path)\n",
    "    os.mkdir(gen_data_path)\n",
    "\n",
    "    train_data.to_csv(\n",
    "        gen_data_path + \"/train_data.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "    test_data.to_csv(\n",
    "        gen_data_path + \"/test_data.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "\n",
    "    train_data_alice = train_data[\n",
    "        [\"user_id\", \"sequence_movie_ids\", \"target_movie_id\", \"label\"]\n",
    "    ]\n",
    "    train_data_bob = train_data[\n",
    "        [\n",
    "            \"gender\",\n",
    "            \"age_group\",\n",
    "            \"occupation\",\n",
    "        ]\n",
    "    ]\n",
    "\n",
    "    test_data_alice = test_data[\n",
    "        [\"user_id\", \"sequence_movie_ids\", \"target_movie_id\", \"label\"]\n",
    "    ]\n",
    "    test_data_bob = test_data[\n",
    "        [\n",
    "            \"gender\",\n",
    "            \"age_group\",\n",
    "            \"occupation\",\n",
    "        ]\n",
    "    ]\n",
    "\n",
    "    train_data_alice.to_csv(\n",
    "        gen_data_path + \"/train_data_alice.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "    train_data_bob.to_csv(\n",
    "        gen_data_path + \"/train_data_bob.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "\n",
    "    test_data_alice.to_csv(\n",
    "        gen_data_path + \"/test_data_alice.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "    test_data_bob.to_csv(\n",
    "        gen_data_path + \"/test_data_bob.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "\n",
    "\n",
    "generate_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6652f32f-3099-4c1a-9b58-dfdedd27e10c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "user_id|sequence_movie_ids|target_movie_id|label\n",
      "0|2374,768,296,28|1435|3.0\n",
      "0|1435,876,2620,228|1950|5.0\n",
      "0|1950,1723,3580,213|3796|4.0\n",
      "0|3796,3460,1825,1014|1089|5.0\n",
      "0|1015,34,1078,539|107|4.0\n",
      "0|107,3791,314,1941|269|4.0\n",
      "0|269,1903,3517,2076|3375|4.0\n",
      "0|3375,2295,1935,35|1414|3.0\n",
      "0|1414,216,3444,1498|582|4.0\n"
     ]
    }
   ],
   "source": [
    "! head ./test_bst_data/data_sl_bst/train_data_alice.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a6598709-0ddb-46cc-8cb2-77891470281d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gender|age_group|occupation\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n",
      "0|0|2\n"
     ]
    }
   ],
   "source": [
    "! head ./test_bst_data/data_sl_bst/train_data_bob.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b144729b-e17a-4676-af63-32274f97de17",
   "metadata": {},
   "source": [
    "到这里，我们以及产出训练测试用的数据文件。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e436048e-0abd-431d-b3b4-4e80b4511a11",
   "metadata": {},
   "source": [
    "## 构造 dataset_builder 处理数据\n",
    "这里对特征进行 ID 化处理，rating > 3 作为正样本，否则作为负样本 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f9419fa-02f3-45a9-b1be-330b877e8fce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "\n",
    "# transform DataFrame to Dataset\n",
    "class AliceDataset(Dataset):\n",
    "    def __init__(self, df, label_df, vocab_dir):\n",
    "        self.df = df\n",
    "        self.label_df = label_df\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        user_id = torch.tensor([int(self.df[\"user_id\"].iloc[index])])\n",
    "        seq_ids = torch.tensor(\n",
    "            [int(sid) for sid in self.df[\"sequence_movie_ids\"].iloc[index].split(',')]\n",
    "        )\n",
    "        target_id = torch.tensor([int(self.df[\"target_movie_id\"].iloc[index])])\n",
    "\n",
    "        label = 1 if self.label_df[\"label\"].iloc[index] > 3 else 0\n",
    "\n",
    "        return (user_id, target_id, seq_ids), label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.label_df)\n",
    "\n",
    "\n",
    "class BobDataset(Dataset):\n",
    "    def __init__(self, df, vocab_dir):\n",
    "        self.df = df\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        gender = torch.tensor([self.df[\"gender\"].iloc[index]])\n",
    "        age = torch.tensor([self.df[\"age_group\"].iloc[index]])\n",
    "        occupation = torch.tensor([self.df[\"occupation\"].iloc[index]])\n",
    "\n",
    "        return (\n",
    "            gender,\n",
    "            age,\n",
    "            occupation,\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "\n",
    "def create_dataset_builder_alice(batch_size=32):\n",
    "    def dataset_builder(x):\n",
    "        data_set = AliceDataset(x[0], x[1], gen_data_path)\n",
    "        dataloader = DataLoader(\n",
    "            dataset=data_set,\n",
    "            batch_size=batch_size,\n",
    "        )\n",
    "        return dataloader\n",
    "\n",
    "    return dataset_builder\n",
    "\n",
    "\n",
    "def create_dataset_builder_bob(batch_size=32):\n",
    "    def dataset_builder(x):\n",
    "        data_set = BobDataset(x[0], gen_data_path)\n",
    "        dataloader = DataLoader(\n",
    "            dataset=data_set,\n",
    "            batch_size=batch_size,\n",
    "        )\n",
    "        return dataloader\n",
    "\n",
    "    return dataset_builder\n",
    "\n",
    "\n",
    "batch_size = 128\n",
    "dataset_buidler_dict = {\n",
    "    alice: create_dataset_builder_alice(batch_size=batch_size),\n",
    "    bob: create_dataset_builder_bob(\n",
    "        batch_size=batch_size,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a731615-01d8-4227-b6b6-f2b7b2894bbc",
   "metadata": {},
   "source": [
    "## 定义模型结构\n",
    "通过 mode 指定 base 模型结构\n",
    "- mode='ori': BSTBase\n",
    "- mode='plus': BSTPlusBase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "95a67d44-443a-4e44-b8b2-36775de21094",
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow_fl.ml.nn.applications.sl_bst_torch import BSTBase, BSTBasePlus, BSTfuse\n",
    "\n",
    "\n",
    "def create_base_model_alice(mode=\"ori\"):\n",
    "    def create_model():\n",
    "        fea_emb_size = {\n",
    "            \"user_id\": [fea_emb_input_size[\"user_id\"], 8],\n",
    "            \"target_id\": [fea_emb_input_size[\"target_id\"], 9],\n",
    "        }\n",
    "\n",
    "        fea_list = [\"user_id\", \"target_id\", \"sequence_movie_ids\"]\n",
    "\n",
    "        if mode == \"ori\":\n",
    "            model = BSTBase(\n",
    "                fea_list=fea_list,\n",
    "                fea_emb_dim=fea_emb_size,\n",
    "                sequence_fea=[\"sequence_movie_ids\"],\n",
    "                target_item_fea=\"target_id\",\n",
    "                seq_len={\"sequence_movie_ids\": 4},\n",
    "                padding_idx=0,\n",
    "            )\n",
    "        else:\n",
    "            model = BSTBasePlus(\n",
    "                fea_list=fea_list,\n",
    "                fea_emb_dim=fea_emb_size,\n",
    "                sequence_fea=[\"sequence_movie_ids\"],\n",
    "                target_item_fea=\"target_id\",\n",
    "                seq_len={\"sequence_movie_ids\": 4},\n",
    "                padding_idx=0,\n",
    "            )\n",
    "\n",
    "        return model\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "def create_base_model_bob():\n",
    "    def create_model():\n",
    "        fea_list = [\"gender\", \"age_group\", \"occupation\"]\n",
    "        fea_emb_size = {}\n",
    "        for key in fea_list:\n",
    "            fea_emb_size[key] = [fea_emb_input_size[key], 8]\n",
    "\n",
    "        model = BSTBase(\n",
    "            fea_list=fea_list,\n",
    "            fea_emb_dim=fea_emb_size,\n",
    "            sequence_fea=[],\n",
    "            target_item_fea=None,\n",
    "            seq_len=None,\n",
    "        )\n",
    "\n",
    "        return model\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "def create_fuse_model(mode=\"ori\"):\n",
    "    def create_model():\n",
    "        if mode == 'ori':\n",
    "            model = BSTfuse(dnn_units_size=[77, 64])\n",
    "        else:\n",
    "            model = BSTfuse(dnn_units_size=[50, 64])\n",
    "        return model\n",
    "\n",
    "    return create_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0f6c821-896e-4507-b9bb-edbb7e47e3a3",
   "metadata": {},
   "source": [
    "## 定义 SL Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "662950ad-18b4-4fb6-bd41-625d347d16e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow_fl.ml.nn.sl.backend.torch.strategy.split_nn.PYUSLTorchModel'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow_fl.ml.nn.sl.backend.torch.strategy.split_nn.PYUSLTorchModel'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "from torch import nn, optim\n",
    "from torchmetrics import AUROC, Accuracy, Precision\n",
    "from secretflow_fl.ml.nn.core.torch import metric_wrapper, optim_wrapper, TorchModel\n",
    "from secretflow_fl.ml.nn import SLModel\n",
    "\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss\n",
    "optim_fn = optim_wrapper(optim.Adam, lr=1e-2)\n",
    "\n",
    "mode = 'plus'\n",
    "base_model_alice = TorchModel(\n",
    "    model_fn=create_base_model_alice(mode),\n",
    "    loss_fn=loss_fn,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(Accuracy, task=\"multiclass\", num_classes=2, average='micro'),\n",
    "        metric_wrapper(Precision, task=\"multiclass\", num_classes=2, average='micro'),\n",
    "        metric_wrapper(AUROC, task=\"binary\"),\n",
    "    ],\n",
    ")\n",
    "\n",
    "base_model_bob = TorchModel(\n",
    "    model_fn=create_base_model_bob(),\n",
    "    loss_fn=loss_fn,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(Accuracy, task=\"multiclass\", num_classes=2, average='micro'),\n",
    "        metric_wrapper(Precision, task=\"multiclass\", num_classes=2, average='micro'),\n",
    "        metric_wrapper(AUROC, task=\"binary\"),\n",
    "    ],\n",
    ")\n",
    "\n",
    "fuse_model = TorchModel(\n",
    "    model_fn=create_fuse_model(mode),\n",
    "    loss_fn=loss_fn,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(Accuracy, task=\"multiclass\", num_classes=2, average='micro'),\n",
    "        metric_wrapper(Precision, task=\"multiclass\", num_classes=2, average='micro'),\n",
    "        metric_wrapper(AUROC, task=\"multiclass\", num_classes=2),\n",
    "    ],\n",
    ")\n",
    "\n",
    "base_model_dict = {\n",
    "    alice: base_model_alice,\n",
    "    bob: base_model_bob,\n",
    "}\n",
    "\n",
    "sl_model = SLModel(\n",
    "    base_model_dict=base_model_dict,\n",
    "    device_y=alice,\n",
    "    model_fuse=fuse_model,\n",
    "    random_seed=1234,\n",
    "    backend='torch',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e8e085e-cc66-45c4-8afa-04331d4bba86",
   "metadata": {},
   "source": [
    "## 开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "46ea8414-e0bf-4e1d-9ba6-0c125c540538",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:SL Train Params: {'x': VDataFrame(partitions={PYURuntime(alice): <secretflow.data.partition.pandas.partition.PdPartition object at 0x7f7529c6d250>, PYURuntime(bob): <secretflow.data.partition.pandas.partition.PdPartition object at 0x7f75274a7760>}, aligned=True), 'y': VDataFrame(partitions={PYURuntime(alice): <secretflow.data.partition.pandas.partition.PdPartition object at 0x7f75274a7580>}, aligned=True), 'batch_size': 128, 'epochs': 1, 'verbose': 1, 'callbacks': None, 'validation_data': (VDataFrame(partitions={PYURuntime(alice): <secretflow.data.partition.pandas.partition.PdPartition object at 0x7f7529c6d250>, PYURuntime(bob): <secretflow.data.partition.pandas.partition.PdPartition object at 0x7f75274a7760>}, aligned=True), VDataFrame(partitions={PYURuntime(alice): <secretflow.data.partition.pandas.partition.PdPartition object at 0x7f75274a7580>}, aligned=True)), 'shuffle': False, 'sample_weight': None, 'validation_freq': 1, 'dp_spent_step_freq': None, 'dataset_builder': {PYURuntime(alice): <function create_dataset_builder_alice.<locals>.dataset_builder at 0x7f7588b27670>, PYURuntime(bob): <function create_dataset_builder_bob.<locals>.dataset_builder at 0x7f7588b27700>}, 'audit_log_params': {}, 'random_seed': 1234, 'audit_log_dir': None, 'self': <secretflow_fl.ml.nn.sl.sl_model.SLModel object at 0x7f7529c6db80>}\n",
      "100%|██████████| 1664/1664 [04:01<00:00,  6.90it/s, epoch: 1/1 -  train_loss:0.6442505717277527  train_MulticlassAccuracy:0.6493232250213623  train_MulticlassPrecision:0.6493232250213623  train_MulticlassAUROC:0.6861191391944885  val_val_loss:0.6395253539085388  val_MulticlassAccuracy:0.6840226650238037  val_MulticlassPrecision:0.6840226650238037  val_MulticlassAUROC:0.734423816204071 ]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "history:  {'train_loss': [array(0.6442506, dtype=float32)], 'train_MulticlassAccuracy': [tensor(0.6493)], 'train_MulticlassPrecision': [tensor(0.6493)], 'train_MulticlassAUROC': [tensor(0.6861)], 'val_val_loss': [array(0.63952535, dtype=float32)], 'val_MulticlassAccuracy': [tensor(0.6840)], 'val_MulticlassPrecision': [tensor(0.6840)], 'val_MulticlassAUROC': [tensor(0.7344)]}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from secretflow.data.vertical import read_csv\n",
    "\n",
    "\n",
    "vdf = read_csv(\n",
    "    {\n",
    "        alice: gen_data_path + '/train_data_alice.csv',\n",
    "        bob: gen_data_path + '/train_data_bob.csv',\n",
    "    },\n",
    "    delimiter='|',\n",
    ")\n",
    "label = vdf[\"label\"]\n",
    "data = vdf.drop(columns=[\"label\"])\n",
    "\n",
    "epoch = 1\n",
    "\n",
    "history = sl_model.fit(\n",
    "    data,\n",
    "    label,\n",
    "    validation_data=(data, label),\n",
    "    epochs=epoch,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    "    random_seed=1234,\n",
    "    dataset_builder=dataset_buidler_dict,\n",
    ")\n",
    "print('history: ', history)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66d28f48-f9f6-414b-be4a-91e6bbe59fe8",
   "metadata": {},
   "source": [
    "## 总结\n",
    "本文通过 movieLens 数据集上的推荐任务来演示了如何通过隐语来训练拆分 BST 模型，您需要\n",
    "1. 下载并处理数据集；\n",
    "2. 构造 dataset_builder 处理数据；\n",
    "3. 定义模型结构，调用 `BSTBase`， `BSTPlusBase` 和 `BSTFuse` 定义行模结构；\n",
    "4. 使用 SLModel 进行训练，预测，评估即可。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0091492a-1b57-4b75-914d-7c03786cc00a",
   "metadata": {},
   "source": [
    "您可以在自己的数据集上进行尝试，如有任何问题，可以在 github 上进行讨论。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
